#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
限流中间件

这个模块包含请求限流中间件：
1. 基于IP的限流
2. 基于用户的限流
3. 基于端点的限流
4. 滑动窗口限流算法
5. 令牌桶限流算法

提供灵活的限流策略
支持自定义限流规则

作者: FastAPI 学习项目
创建时间: 2024
"""

import time
import asyncio
from typing import Callable, Dict, Optional, List, Tuple
from fastapi import Request, Response, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from loguru import logger
from collections import defaultdict, deque
from datetime import datetime, timedelta
import threading
import hashlib
import json
from enum import Enum


class RateLimitStrategy(Enum):
    """
    限流策略枚举
    """
    FIXED_WINDOW = "fixed_window"  # 固定窗口
    SLIDING_WINDOW = "sliding_window"  # 滑动窗口
    TOKEN_BUCKET = "token_bucket"  # 令牌桶
    LEAKY_BUCKET = "leaky_bucket"  # 漏桶


class RateLimitRule:
    """
    限流规则类
    
    定义单个限流规则的配置。
    """
    
    def __init__(
        self,
        *,
        requests: int,
        window: int,
        strategy: RateLimitStrategy = RateLimitStrategy.SLIDING_WINDOW,
        burst: Optional[int] = None,
        key_func: Optional[Callable] = None,
        exempt_ips: List[str] = None,
        exempt_paths: List[str] = None
    ):
        """
        初始化限流规则
        
        Args:
            requests: 允许的请求数量
            window: 时间窗口（秒）
            strategy: 限流策略
            burst: 突发请求数量（仅用于令牌桶）
            key_func: 自定义键生成函数
            exempt_ips: 豁免的IP列表
            exempt_paths: 豁免的路径列表
        """
        self.requests = requests
        self.window = window
        self.strategy = strategy
        self.burst = burst or requests
        self.key_func = key_func
        self.exempt_ips = exempt_ips or []
        self.exempt_paths = exempt_paths or []
    
    def should_exempt(self, request: Request, client_ip: str) -> bool:
        """
        判断请求是否应该豁免限流
        
        Args:
            request: HTTP请求对象
            client_ip: 客户端IP
            
        Returns:
            bool: 是否豁免
        """
        # 检查IP豁免
        if client_ip in self.exempt_ips:
            return True
        
        # 检查路径豁免
        path = request.url.path
        for exempt_path in self.exempt_paths:
            if path.startswith(exempt_path):
                return True
        
        return False
    
    def generate_key(self, request: Request, client_ip: str) -> str:
        """
        生成限流键
        
        Args:
            request: HTTP请求对象
            client_ip: 客户端IP
            
        Returns:
            str: 限流键
        """
        if self.key_func:
            return self.key_func(request, client_ip)
        
        # 默认使用IP作为键
        return f"rate_limit:{client_ip}"


class SlidingWindowCounter:
    """
    滑动窗口计数器
    
    实现滑动窗口限流算法。
    """
    
    def __init__(self, window_size: int, max_requests: int):
        """
        初始化滑动窗口计数器
        
        Args:
            window_size: 窗口大小（秒）
            max_requests: 最大请求数
        """
        self.window_size = window_size
        self.max_requests = max_requests
        self.requests = deque()
        self._lock = threading.Lock()
    
    def is_allowed(self) -> Tuple[bool, int, int]:
        """
        检查是否允许请求
        
        Returns:
            Tuple[bool, int, int]: (是否允许, 剩余请求数, 重置时间)
        """
        with self._lock:
            current_time = time.time()
            
            # 移除过期的请求记录
            while self.requests and self.requests[0] <= current_time - self.window_size:
                self.requests.popleft()
            
            # 检查是否超过限制
            if len(self.requests) >= self.max_requests:
                # 计算重置时间
                reset_time = int(self.requests[0] + self.window_size)
                return False, 0, reset_time
            
            # 添加当前请求
            self.requests.append(current_time)
            
            # 计算剩余请求数和重置时间
            remaining = self.max_requests - len(self.requests)
            reset_time = int(current_time + self.window_size)
            
            return True, remaining, reset_time


class TokenBucket:
    """
    令牌桶
    
    实现令牌桶限流算法。
    """
    
    def __init__(self, capacity: int, refill_rate: float, refill_period: float = 1.0):
        """
        初始化令牌桶
        
        Args:
            capacity: 桶容量
            refill_rate: 填充速率（每个周期添加的令牌数）
            refill_period: 填充周期（秒）
        """
        self.capacity = capacity
        self.refill_rate = refill_rate
        self.refill_period = refill_period
        self.tokens = capacity
        self.last_refill = time.time()
        self._lock = threading.Lock()
    
    def is_allowed(self, tokens_requested: int = 1) -> Tuple[bool, int, int]:
        """
        检查是否允许请求
        
        Args:
            tokens_requested: 请求的令牌数
            
        Returns:
            Tuple[bool, int, int]: (是否允许, 剩余令牌数, 重置时间)
        """
        with self._lock:
            current_time = time.time()
            
            # 计算需要添加的令牌数
            time_passed = current_time - self.last_refill
            tokens_to_add = (time_passed / self.refill_period) * self.refill_rate
            
            # 更新令牌数（不超过容量）
            self.tokens = min(self.capacity, self.tokens + tokens_to_add)
            self.last_refill = current_time
            
            # 检查是否有足够的令牌
            if self.tokens >= tokens_requested:
                self.tokens -= tokens_requested
                reset_time = int(current_time + (self.capacity - self.tokens) / self.refill_rate * self.refill_period)
                return True, int(self.tokens), reset_time
            else:
                # 计算何时会有足够的令牌
                tokens_needed = tokens_requested - self.tokens
                wait_time = (tokens_needed / self.refill_rate) * self.refill_period
                reset_time = int(current_time + wait_time)
                return False, int(self.tokens), reset_time


class FixedWindowCounter:
    """
    固定窗口计数器
    
    实现固定窗口限流算法。
    """
    
    def __init__(self, window_size: int, max_requests: int):
        """
        初始化固定窗口计数器
        
        Args:
            window_size: 窗口大小（秒）
            max_requests: 最大请求数
        """
        self.window_size = window_size
        self.max_requests = max_requests
        self.count = 0
        self.window_start = time.time()
        self._lock = threading.Lock()
    
    def is_allowed(self) -> Tuple[bool, int, int]:
        """
        检查是否允许请求
        
        Returns:
            Tuple[bool, int, int]: (是否允许, 剩余请求数, 重置时间)
        """
        with self._lock:
            current_time = time.time()
            
            # 检查是否需要重置窗口
            if current_time - self.window_start >= self.window_size:
                self.count = 0
                self.window_start = current_time
            
            # 检查是否超过限制
            if self.count >= self.max_requests:
                reset_time = int(self.window_start + self.window_size)
                return False, 0, reset_time
            
            # 增加计数
            self.count += 1
            remaining = self.max_requests - self.count
            reset_time = int(self.window_start + self.window_size)
            
            return True, remaining, reset_time


class RateLimitStorage:
    """
    限流存储
    
    管理所有限流计数器的存储。
    """
    
    def __init__(self):
        """
        初始化限流存储
        """
        self.counters = {}
        self._lock = threading.Lock()
        
        # 启动清理任务
        self._start_cleanup_task()
    
    def get_counter(self, key: str, rule: RateLimitRule):
        """
        获取或创建计数器
        
        Args:
            key: 限流键
            rule: 限流规则
            
        Returns:
            计数器对象
        """
        with self._lock:
            if key not in self.counters:
                if rule.strategy == RateLimitStrategy.SLIDING_WINDOW:
                    self.counters[key] = {
                        'counter': SlidingWindowCounter(rule.window, rule.requests),
                        'last_access': time.time()
                    }
                elif rule.strategy == RateLimitStrategy.TOKEN_BUCKET:
                    refill_rate = rule.requests / rule.window
                    self.counters[key] = {
                        'counter': TokenBucket(rule.burst, refill_rate, 1.0),
                        'last_access': time.time()
                    }
                elif rule.strategy == RateLimitStrategy.FIXED_WINDOW:
                    self.counters[key] = {
                        'counter': FixedWindowCounter(rule.window, rule.requests),
                        'last_access': time.time()
                    }
                else:
                    # 默认使用滑动窗口
                    self.counters[key] = {
                        'counter': SlidingWindowCounter(rule.window, rule.requests),
                        'last_access': time.time()
                    }
            
            # 更新最后访问时间
            self.counters[key]['last_access'] = time.time()
            return self.counters[key]['counter']
    
    def cleanup_expired_counters(self, max_idle_time: int = 3600):
        """
        清理过期的计数器
        
        Args:
            max_idle_time: 最大空闲时间（秒）
        """
        with self._lock:
            current_time = time.time()
            expired_keys = []
            
            for key, data in self.counters.items():
                if current_time - data['last_access'] > max_idle_time:
                    expired_keys.append(key)
            
            for key in expired_keys:
                del self.counters[key]
            
            if expired_keys:
                logger.info(f"清理了 {len(expired_keys)} 个过期的限流计数器")
    
    def _start_cleanup_task(self):
        """
        启动清理任务
        """
        def cleanup_loop():
            while True:
                try:
                    time.sleep(300)  # 每5分钟清理一次
                    self.cleanup_expired_counters()
                except Exception as e:
                    logger.error(f"限流计数器清理异常: {e}")
        
        cleanup_thread = threading.Thread(target=cleanup_loop, daemon=True)
        cleanup_thread.start()
    
    def get_statistics(self) -> Dict:
        """
        获取存储统计信息
        
        Returns:
            Dict: 统计信息
        """
        with self._lock:
            return {
                'total_counters': len(self.counters),
                'counter_keys': list(self.counters.keys())[:10]  # 只显示前10个键
            }


class RateLimitMiddleware(BaseHTTPMiddleware):
    """
    限流中间件
    
    提供灵活的请求限流功能。
    """
    
    def __init__(
        self,
        app,
        *,
        default_rule: RateLimitRule = None,
        rules: Dict[str, RateLimitRule] = None,
        enable_headers: bool = True,
        storage: RateLimitStorage = None
    ):
        """
        初始化限流中间件
        
        Args:
            app: FastAPI 应用实例
            default_rule: 默认限流规则
            rules: 特定路径的限流规则
            enable_headers: 是否添加限流头部
            storage: 自定义存储
        """
        super().__init__(app)
        
        # 默认规则：每分钟100个请求
        self.default_rule = default_rule or RateLimitRule(
            requests=100,
            window=60,
            strategy=RateLimitStrategy.SLIDING_WINDOW
        )
        
        self.rules = rules or {}
        self.enable_headers = enable_headers
        self.storage = storage or RateLimitStorage()
        
        # 统计信息
        self.stats = {
            'total_requests': 0,
            'blocked_requests': 0,
            'last_reset': time.time()
        }
        self._stats_lock = threading.Lock()
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """
        处理请求并执行限流检查
        
        Args:
            request: HTTP请求对象
            call_next: 下一个中间件或路由处理器
            
        Returns:
            Response: HTTP响应对象
        """
        client_ip = self._get_client_ip(request)
        
        # 更新统计信息
        with self._stats_lock:
            self.stats['total_requests'] += 1
        
        # 获取适用的规则
        rule = self._get_applicable_rule(request)
        
        # 检查是否豁免
        if rule.should_exempt(request, client_ip):
            return await call_next(request)
        
        # 执行限流检查
        allowed, remaining, reset_time = await self._check_rate_limit(request, client_ip, rule)
        
        if not allowed:
            # 更新统计信息
            with self._stats_lock:
                self.stats['blocked_requests'] += 1
            
            # 记录限流事件
            logger.warning(
                "请求被限流",
                extra={
                    "event_type": "rate_limit_exceeded",
                    "client_ip": client_ip,
                    "path": request.url.path,
                    "method": request.method,
                    "rule": {
                        "requests": rule.requests,
                        "window": rule.window,
                        "strategy": rule.strategy.value
                    },
                    "reset_time": reset_time
                }
            )
            
            # 创建限流响应
            response = Response(
                content=json.dumps({
                    "error": "Rate limit exceeded",
                    "message": f"请求过于频繁，请在 {reset_time - int(time.time())} 秒后重试",
                    "retry_after": reset_time - int(time.time())
                }),
                status_code=429,
                media_type="application/json"
            )
            
            # 添加限流头部
            if self.enable_headers:
                self._add_rate_limit_headers(response, rule, remaining, reset_time)
            
            return response
        
        # 处理请求
        response = await call_next(request)
        
        # 添加限流头部
        if self.enable_headers:
            self._add_rate_limit_headers(response, rule, remaining, reset_time)
        
        return response
    
    def _get_client_ip(self, request: Request) -> str:
        """
        获取客户端IP地址
        
        Args:
            request: HTTP请求对象
            
        Returns:
            str: 客户端IP地址
        """
        # 检查代理头部
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip
        
        return request.client.host if request.client else "unknown"
    
    def _get_applicable_rule(self, request: Request) -> RateLimitRule:
        """
        获取适用的限流规则
        
        Args:
            request: HTTP请求对象
            
        Returns:
            RateLimitRule: 适用的规则
        """
        path = request.url.path
        
        # 检查特定路径规则
        for pattern, rule in self.rules.items():
            if path.startswith(pattern):
                return rule
        
        # 返回默认规则
        return self.default_rule
    
    async def _check_rate_limit(
        self, 
        request: Request, 
        client_ip: str, 
        rule: RateLimitRule
    ) -> Tuple[bool, int, int]:
        """
        检查限流
        
        Args:
            request: HTTP请求对象
            client_ip: 客户端IP
            rule: 限流规则
            
        Returns:
            Tuple[bool, int, int]: (是否允许, 剩余请求数, 重置时间)
        """
        # 生成限流键
        key = rule.generate_key(request, client_ip)
        
        # 获取计数器
        counter = self.storage.get_counter(key, rule)
        
        # 检查是否允许
        return counter.is_allowed()
    
    def _add_rate_limit_headers(
        self, 
        response: Response, 
        rule: RateLimitRule, 
        remaining: int, 
        reset_time: int
    ):
        """
        添加限流头部
        
        Args:
            response: HTTP响应对象
            rule: 限流规则
            remaining: 剩余请求数
            reset_time: 重置时间
        """
        response.headers["X-RateLimit-Limit"] = str(rule.requests)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        response.headers["X-RateLimit-Reset"] = str(reset_time)
        response.headers["X-RateLimit-Window"] = str(rule.window)
        response.headers["X-RateLimit-Strategy"] = rule.strategy.value
    
    def get_statistics(self) -> Dict:
        """
        获取限流统计信息
        
        Returns:
            Dict: 统计信息
        """
        with self._stats_lock:
            uptime = time.time() - self.stats['last_reset']
            return {
                'total_requests': self.stats['total_requests'],
                'blocked_requests': self.stats['blocked_requests'],
                'block_rate': (
                    self.stats['blocked_requests'] / self.stats['total_requests'] * 100
                    if self.stats['total_requests'] > 0 else 0
                ),
                'uptime_seconds': int(uptime),
                'requests_per_second': (
                    self.stats['total_requests'] / uptime
                    if uptime > 0 else 0
                ),
                'storage_stats': self.storage.get_statistics()
            }
    
    def reset_statistics(self):
        """
        重置统计信息
        """
        with self._stats_lock:
            self.stats = {
                'total_requests': 0,
                'blocked_requests': 0,
                'last_reset': time.time()
            }
        logger.info("限流统计信息已重置")
    
    def add_rule(self, pattern: str, rule: RateLimitRule):
        """
        添加限流规则
        
        Args:
            pattern: 路径模式
            rule: 限流规则
        """
        self.rules[pattern] = rule
        logger.info(f"添加限流规则: {pattern} -> {rule.requests}/{rule.window}s")
    
    def remove_rule(self, pattern: str):
        """
        移除限流规则
        
        Args:
            pattern: 路径模式
        """
        if pattern in self.rules:
            del self.rules[pattern]
            logger.info(f"移除限流规则: {pattern}")